package com.xiashitech.agent.instrumentation.db.common;

import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import lombok.*;

import java.util.*;


public class XSResultSizeContext {

    public static Long getSize(Object resultSet) {
        ResultSize resultSizeForRecord = resultSize.get().get(resultSet);
        if(resultSizeForRecord != null) {
            return resultSizeForRecord.getSize();
        } else {
            return 0L;
        }
    }

    public static Boolean getIsRecord(Object resultSet) {
        ResultSize resultSizeForRecord = resultSize.get().get(resultSet);
        if(resultSizeForRecord != null) {
            return resultSizeForRecord.getIsRecord();
        } else {
            return false;
        }
    }

    public static void setIsRecord(Object resultSet,Boolean isRecord) {
        ResultSize resultSizeForRecord = resultSize.get().computeIfAbsent(resultSet, resultSetKey -> ResultSize.builder().build());
        if(resultSizeForRecord != null) {
            resultSizeForRecord.setIsRecord(isRecord);
        }
    }

    public static void record(Object resultSet) {
        ResultSize resultSizeForRecord = XSResultSizeContext.resultSize.get().computeIfAbsent(resultSet, resultSetKey -> ResultSize.builder().build());
        if(resultSizeForRecord != null && resultSizeForRecord.getIsRecord()) {
            resultSizeForRecord.increase();
        }
    }

    public static ThreadLocal<Map<Object,ResultSize>> resultSize = new ThreadLocal<Map<Object,ResultSize>>() {
        protected Map<Object,ResultSize> initialValue() {
            return new HashMap<Object,ResultSize>();
        }
    };

    public static ResultSetContext getResultSetContext(Object resultSet) {
        Optional<ResultSetContext> resultSetContextOptional = XSResultSizeContext.resultSetContexts.get().stream().filter(context -> resultSet == context.getResultSet()).findFirst();
        if (resultSetContextOptional.isPresent()) {
            return resultSetContextOptional.get();
        } else {
            return null;
        }
    }

    public static void addResultSetContext(ResultSetContext resultSetContext) {
        Optional<ResultSetContext> resultSetContextOptional = XSResultSizeContext.resultSetContexts.get().stream().filter(context -> context.getResultSet() == resultSetContext.getResultSet()).findFirst();
        if(!resultSetContextOptional.isPresent()) {
            XSResultSizeContext.resultSetContexts.get().add(resultSetContext);
        }
    }

    public static ThreadLocal<List<ResultSetContext>> resultSetContexts = new ThreadLocal<List<ResultSetContext>>() {
        protected List<ResultSetContext> initialValue() {
            return new ArrayList<ResultSetContext>();
        }
    };

    public static StatementSQLContext getStatementSQLContext(Object statement) {
        Optional<StatementSQLContext> statementSQLContextOptional = statementSQLContexts.get().stream().filter(context -> statement == context.getStatement()).findFirst();
        if (statementSQLContextOptional.isPresent()) {
            return statementSQLContextOptional.get();
        } else {
            return null;
        }
    }

    public static void addStatementSQLContext(StatementSQLContext statementSQLContext) {
        Optional<StatementSQLContext> statementSQLContextOptional = statementSQLContexts.get().stream().filter(context -> context.getStatement() == statementSQLContext.getStatement()).findFirst();
        if(!statementSQLContextOptional.isPresent()) {
            statementSQLContexts.get().add(statementSQLContext);
        }
    }

    public static ThreadLocal<List<StatementSQLContext>> statementSQLContexts = new ThreadLocal<List<StatementSQLContext>>() {
        protected List<StatementSQLContext> initialValue() {
            return new ArrayList<StatementSQLContext>();
        }
    };

    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class Request {
        Object resultSet;
        String sql;
        Long size;
    }

    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class ResultSetContext {
        Context context;
        Scope scope;
        Object resultSet;
        Throwable throwable;
        String sql;
    }

    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class ResultSize {
        @Builder.Default
        Boolean isRecord = false;
        @Builder.Default
        Long size = 0L;
        public void increase() {
            size++;
        }
    }

    @Getter
    @Setter
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class StatementSQLContext {
        Object statement;
        String sql;
    }

    public static void refreshCurrentThreadLocal() {
        resultSize.set(new HashMap<Object,ResultSize>());
        resultSetContexts.set(new ArrayList<ResultSetContext>());
        statementSQLContexts.set(new ArrayList<StatementSQLContext>());
    }
}
